package cn.xiaopengstack.domain.agent.service.armory.node;

import cn.bugstack.wrench.design.framework.tree.StrategyHandler;
import cn.xiaopengstack.domain.agent.model.entity.ArmoryCommandEntity;
import cn.xiaopengstack.domain.agent.model.valobj.*;
import cn.xiaopengstack.domain.agent.service.armory.AbstractArmorySupport;
import cn.xiaopengstack.domain.agent.service.armory.factory.DefaultArmoryStrategyFactory;
import com.alibaba.fastjson.JSONObject;
import io.modelcontextprotocol.client.McpSyncClient;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.BeansException;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author jiangyangang
 * AI对话模型工具组装注册节点
 */
@Slf4j
@Component
public class AiClientModelNode extends AbstractArmorySupport {

    @Resource
    private AiClientAdvisorNode aiClientAdvisorNode;

    @Override
    protected String beanName(String id) {
        return AiAgentEnumVO.AI_CLIENT_MODEL.getBeanName(id);
    }

    @Override
    protected String dataName() {
        return AiAgentEnumVO.AI_CLIENT_MODEL.getDataName();
    }

    @Override
    protected String doApply(ArmoryCommandEntity requestParameter, DefaultArmoryStrategyFactory.DynamicContext dynamicContext) throws Exception {


        List<AiClientModelVO> modelVOS = dynamicContext.getValue(dataName());
        if (CollectionUtils.isEmpty(modelVOS)) {
            log.info("AI客户端加载，无可加载的Model...");
            return router(requestParameter, dynamicContext);
        }

        log.info("AI客户端加载，Model加载中【{}】", JSONObject.toJSONString(modelVOS.stream().map(model -> model.getModelId() + "_" + model.getModelName()).collect(Collectors.joining(","))));

        for (AiClientModelVO modelVO : modelVOS) {
            List<AiClientToolMcpVO> mcpVOS = (List<AiClientToolMcpVO>) Optional.of(dynamicContext.getValue(AiAgentEnumVO.AI_CLIENT_TOOL_MCP.getDataName())).orElse(List.of());

            List<AiClientToolMcpVO> mcpFilterList = mcpVOS.stream().filter(mcp -> modelVO.getToolMcpIds().contains(mcp.getMcpId())).toList();

            OpenAiApi openAiApi = null;
            try {
                openAiApi = (OpenAiApi) applicationContext.getBean(AiAgentEnumVO.AI_CLIENT_API.getBeanNameTag() + modelVO.getApiId());
            } catch (BeansException e) {
                log.warn("数据装载，AI对话模型客户端节点，无法获取OpenAiApi，按预期开始重新初始化!");
            }

            if (Objects.isNull(openAiApi)) {
                // 重新构造openAiApi
                List<AiClientApiVO> apiVOS = dynamicContext.getValue(AiAgentEnumVO.AI_CLIENT_API.getDataName());
                AiClientApiVO aiClientApiVO = apiVOS.stream().filter(api -> api.getApiId().equals(modelVO.getApiId())).findFirst().get();
                openAiApi = OpenAiApi.builder()
                        .apiKey(aiClientApiVO.getApiKey())
                        .baseUrl(aiClientApiVO.getBaseUrl())
                        .embeddingsPath(aiClientApiVO.getEmbeddingsPath())
                        .completionsPath(aiClientApiVO.getCompletionsPath())
                        .build();
                registerBean(AiAgentEnumVO.AI_CLIENT_API.getBeanNameTag() + modelVO.getApiId(), OpenAiApi.class, openAiApi);
            }

            OpenAiChatModel openAiChatModel = OpenAiChatModel.builder()
                    .openAiApi(openAiApi)
                    .defaultOptions(OpenAiChatOptions.builder()
                            .model(modelVO.getModelName())
                            // mcpClient移至chatClient进行注入
//                            .toolCallbacks(getMcpClients(mcpFilterList))
                            .build())
                    .build();

            registerBean(beanName(modelVO.getModelId()), OpenAiChatModel.class, openAiChatModel);
        }

        return router(requestParameter, dynamicContext);
    }

    private List<ToolCallback> getMcpClients(List<AiClientToolMcpVO> mcpFilterList) {
        List<McpSyncClient> mcpSyncClients = mcpFilterList.stream().map(this::getMcpClient).toList();
        return Arrays.asList(new SyncMcpToolCallbackProvider(mcpSyncClients).getToolCallbacks());
    }

    @Override
    public StrategyHandler<ArmoryCommandEntity, DefaultArmoryStrategyFactory.DynamicContext, String> get(ArmoryCommandEntity requestParameter, DefaultArmoryStrategyFactory.DynamicContext dynamicContext) throws Exception {
        return aiClientAdvisorNode;
    }
}
